package mr

import "fmt"
import "log"
import "net"
import "os"
import "sync"
// import "errors"
import "net/rpc"
import "net/http"
import "path/filepath"
import "time"

type Master struct {
	// Your definitions here.

}

// Your code here -- RPC handlers for the worker to call.

type TaskMap struct {
	FileName string
	IsAlloc bool
	IsFinish bool
	MapTi time.Time
}

type TaskReduce struct {
	ReduceIdx int
	IsAlloc bool
	IsFinish bool
	MapTi time.Time
}

type TaskAssign struct {
	mapTasks []*TaskMap
	reduceTasks []*TaskReduce
	nReduce int
	mux sync.Mutex
}

var tasks = TaskAssign{}

func (t *TaskAssign) InitTasks(nReduce int, files []string) {
	t.mux.Lock()
	t.nReduce = nReduce
	t.mapTasks = make([]*TaskMap, len(files))
	for idx, fi := range files {
		t.mapTasks[idx] = &TaskMap{fi, false, false, time.Now()}
	}
	t.reduceTasks = make([]*TaskReduce, nReduce)
	for i := 0; i < nReduce; i++ {
		t.reduceTasks[i] = &TaskReduce{ReduceIdx: i}
	}
	t.mux.Unlock()
}

func (t *TaskAssign) GetAvaliableMapTask() (int, string) {
	t.mux.Lock()
	defer t.mux.Unlock()
	
	ten := 10 * time.Second
	now := time.Now()
	// find avaliable tasks
	for idx, task := range t.mapTasks {
		if task.IsFinish {
			continue
		}
		// check expire
		if task.IsAlloc {
			expire := now.Sub(task.MapTi)
			if expire > ten {
				task.IsAlloc = false
			}
		}

		if !task.IsAlloc {
			task.IsAlloc = true
			task.MapTi = time.Now()
			return idx, task.FileName
		}
	}
	
	return -1, "nil"
}

func (t *TaskAssign) GetAvaliableReduceTask() int {
	t.mux.Lock()
	defer t.mux.Unlock()
	
	ten := 10 * time.Second
	now := time.Now()
	// find avaliable tasks
	for _, task := range t.reduceTasks {
		if task.IsFinish {
			continue
		}
		// check expire
		if task.IsAlloc {
			expire := now.Sub(task.MapTi)
			if expire > ten {
				task.IsAlloc = false
			}
		}

		if !task.IsAlloc {
			task.IsAlloc = true
			task.MapTi = time.Now()
			return task.ReduceIdx
		}
	}
	
	return -1
}

func (t *TaskAssign) MapFinish() bool {
	t.mux.Lock()
	defer t.mux.Unlock()
	map_finish := true
	for _, t := range tasks.mapTasks {
		map_finish = map_finish && t.IsFinish
	}
	return map_finish
}

func (t *TaskAssign) IsFinish() bool {
	t.mux.Lock()
	defer t.mux.Unlock()
	map_finish := true
	for _, t := range tasks.mapTasks {
		map_finish = map_finish && t.IsFinish
	}
	for _, t := range tasks.reduceTasks {
		map_finish = map_finish && t.IsFinish
	}
	return map_finish
}

func (t *TaskAssign) FinishMapTask(mpIdx int) {
	t.mux.Lock()
	defer t.mux.Unlock()
	t.mapTasks[mpIdx].IsFinish = true
}

func (t *TaskAssign) FinishReduceTask(reduceIdx int) {
	t.mux.Lock()
	defer t.mux.Unlock()
	t.reduceTasks[reduceIdx].IsFinish = true
}


func (m *Master) HandleTask(args *TaskArgs, reply *TaskReply) error {
	reply.WorkType = TASK_WAIT
	reply.NMap = len(tasks.mapTasks)
	switch args.MsgType {
	case TASK_APPLY:
		m.assignMapTask(reply)
		if tasks.MapFinish() { // start reduce task
			m.assignReduceTask(reply)
		}
	case MAP_FINISH:
		tasks.FinishMapTask(args.TaskIdx)
	case REDUCE_FINISH:
		tasks.FinishReduceTask(args.TaskIdx)
	default:
	}
	if tasks.IsFinish() {
		reply.WorkType = WORK_FINISH
		// log.Println("Task Finish!!!")
	}
	// fmt.Printf("HandleMapTask: %v %v\n", args, reply)
	return nil
}

func (m *Master) assignMapTask(reply *TaskReply) bool {
	idx, filename := tasks.GetAvaliableMapTask()
	// log.Printf("Assign map task: %v %v\n", idx, filename)
	if idx != -1 {
		reply.MapIdx = idx
		reply.NReduce = tasks.nReduce
		reply.WorkType = WORK_MAP
		reply.FileName = filename
		return true
	}
	return false
}

func (m *Master) assignReduceTask(reply *TaskReply) bool {
	reduceIdx := tasks.GetAvaliableReduceTask()
	// log.Printf("Assign reduce task: %v\n", reduceIdx)
	if reduceIdx != -1 {
		reply.NReduce = tasks.nReduce
		reply.WorkType = WORK_REDUCE
		reply.ReduceIdx = reduceIdx
		return true
	}
	return false
}

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (m *Master) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}


//
// start a thread that listens for RPCs from worker.go
//
func (m *Master) server() {
	rpc.Register(m)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := masterSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

//
// main/mrmaster.go calls Done() periodically to find out
// if the entire job has finished.
//
func (m *Master) Done() bool {
	// Your code here.

	return tasks.IsFinish()
}

func describe(i interface{}) {
	fmt.Printf("(%v, %T)\n", i, i)
}

// remove all intermedia files of previous tasks
func removeIntermediaFiles() {
	files, _ := filepath.Glob("mr-[0-9]*-[0-9]*")
	for _, fi := range files {
		os.Remove(fi)
	}
}

//
// create a Master.
// main/mrmaster.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeMaster(files []string, nReduce int) *Master {
	m := Master{}

	// Your code here.
	removeIntermediaFiles()

	tasks.InitTasks(nReduce, files)

	// describe(&TaskMap)
	m.server()
	return &m
}
